{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Pocky Query Tool: Automated CVE PoC Search & Validation\n",
        "\n",
        "A lightweight, web-scale agent that helps you find, filter, and fetch real-world PoC exploits — so you don't have to.\n",
        "\n",
        "**Features:**\n",
        "- Automatically searches multiple security-related websites\n",
        "- Intelligently analyzes and extracts PoC code\n",
        "- Automatically selects the most reliable PoC samples\n",
        "- Supports collection of PoCs from multiple sources"
      ],
      "metadata": {
        "id": "BdX56iM1r5aB"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/DhivyaBharathy-web/PraisonAI/blob/main/examples/cookbooks/Pocky_Cybersecurity_PoC_Agent.ipynb)\n"
      ],
      "metadata": {
        "id": "VW766102tUUY"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Install Dependencies"
      ],
      "metadata": {
        "id": "cBu2iXmJsVqE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install praisonaiagents exa-py python-dotenv requests beautifulsoup4"
      ],
      "metadata": {
        "id": "VvbA3A7XsTFm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Set API Keys"
      ],
      "metadata": {
        "id": "Q6B2VEkFsB32"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "# Set your API keys here (replace with your actual keys)\n",
        "os.environ[\"EXA_API_KEY\"] = \"your api key\"\n",
        "os.environ[\"OPENAI_API_KEY\"] = \"your api key\"\n",
        "os.environ[\"OPENAI_BASE_URL\"] = \"https://api.openai.com/v1\"  # Optional, for custom OpenAI endpoints"
      ],
      "metadata": {
        "id": "OlOI3yc_sAkN"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Tools (Core Functions)"
      ],
      "metadata": {
        "id": "1BvjjjJdsanO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "from openai import OpenAI\n",
        "from exa_py import Exa\n",
        "\n",
        "# Dummy/Minimal agent classes for notebook demo\n",
        "class ValidationAgent:\n",
        "    def __init__(self, input_json):\n",
        "        self.input_json = input_json\n",
        "    def run(self):\n",
        "        # Dummy validation logic for notebook demo\n",
        "        data = json.loads(self.input_json)\n",
        "        return True if \"attack_intent\" in data and \"poc_sample\" in data else False\n",
        "\n",
        "class AttackIntentAgent:\n",
        "    def __init__(self, description):\n",
        "        self.description = description\n",
        "    def run(self):\n",
        "        # Dummy intent extraction for notebook demo\n",
        "        return f\"Intent for: {self.description[:50]}...\""
      ],
      "metadata": {
        "id": "GYfAJLXOsbga"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## YAML Prompt (Validation Example)\n",
        "This is the prompt used for PoC validation."
      ],
      "metadata": {
        "id": "THrET8-psjx-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "validation_prompt = \"\"\"\n",
        "You are a highly skilled technical assistant with deep expertise in PoC sample validation.\n",
        "\n",
        "Given the attack intent of a CVE vulnerability and a PoC sample gathered from public sources, your task is to analyze whether the PoC correctly implements the intended attack behavior.\n",
        "\n",
        "Specifically:\n",
        "- Understand the CVE's attack intent, including the attack goal and the underlying exploitation mechanism.\n",
        "- Analyze the PoC to determine whether it is designed to achieve this intent.\n",
        "- Check whether the payloads, request structures, and overall logic of the PoC align with the described attack intent.\n",
        "- You do not need to execute the PoC. Focus on static validation through reasoning and consistency.\n",
        "\n",
        "Your output must be a JSON object with two fields:\n",
        "- \"valid\": a boolean indicating whether the PoC correctly reflects the attack intent.\n",
        "- \"reasoning\": a brief explanation of your judgment. If \"valid\" is false, the reasoning must clearly explain what is incorrect or inconsistent in the PoC compared to the attack intent, so that the PoC can be revised accordingly.\n",
        "\"\"\"\n",
        "print(validation_prompt)"
      ],
      "metadata": {
        "id": "9q3aKl1xshrb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Main (Query and Validate a CVE PoC)"
      ],
      "metadata": {
        "id": "531EZgcLsqP6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def run_pocky_for_cve(cve_id):\n",
        "    # Example: Simulate fetching a description and PoC (replace with real logic)\n",
        "    description = f\"Description for {cve_id} (replace with real Exa/OpenAI search)\"\n",
        "    poc_sample = f\"PoC code for {cve_id} (replace with real PoC search)\"\n",
        "\n",
        "    # Stage 2: Attack Intent\n",
        "    intent = AttackIntentAgent(description).run()\n",
        "    print(f\"Attack Intent: {intent}\")\n",
        "\n",
        "    # Stage 3: Validation\n",
        "    validation_input = json.dumps({\"attack_intent\": intent, \"poc_sample\": poc_sample}, indent=2)\n",
        "    valid = ValidationAgent(validation_input).run()\n",
        "    print(f\"Validation Result: {valid}\")\n",
        "    if valid:\n",
        "        print(f\"PoC for {cve_id} is valid and ready to use.\")\n",
        "    else:\n",
        "        print(f\"PoC for {cve_id} failed validation.\")"
      ],
      "metadata": {
        "id": "PQvtF-RqsrP6"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Example Usage"
      ],
      "metadata": {
        "id": "XiQOiSz8su3u"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "run_pocky_for_cve(\"CVE-2023-4450\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "39w-qpecswjw",
        "outputId": "cdcb3b29-7338-4e3e-b160-5f9568c194ca"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Attack Intent: Intent for: Description for CVE-2023-4450 (replace with real E...\n",
            "Validation Result: True\n",
            "PoC for CVE-2023-4450 is valid and ready to use.\n"
          ]
        }
      ]
    }
  ]
}
